{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# {meth}`torch.Tensor.expand` 转换"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/media/pc/data/lxw/ai/tvm-book/doc/tutorials/frontend\n"
     ]
    }
   ],
   "source": [
    "%cd ../../..\n",
    "import set_env\n",
    "from d2py.utils.file import mkdir\n",
    "temp_dir = \".temp\"\n",
    "mkdir(temp_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Exported graph: graph(%data : Float(*, 7, strides=[7, 1], requires_grad=0, device=cpu)):\n",
      "  %/Constant_output_0 : Long(4, strides=[1], device=cpu) = onnx::Constant[value= 1  3 -1  7 [ CPULongType{4} ], onnx_name=\"/Constant\"](), scope: __main__.Model:: # /tmp/ipykernel_393401/573755285.py:6:0\n",
      "  %/Constant_1_output_0 : Long(1, strides=[1], requires_grad=0, device=cpu) = onnx::Constant[value={4}, onnx_name=\"/Constant_1\"](), scope: __main__.Model:: # /tmp/ipykernel_393401/573755285.py:6:0\n",
      "  %/ConstantOfShape_output_0 : Long(4, strides=[1], device=cpu) = onnx::ConstantOfShape[value={1}, onnx_name=\"/ConstantOfShape\"](%/Constant_1_output_0), scope: __main__.Model:: # /tmp/ipykernel_393401/573755285.py:6:0\n",
      "  %/Constant_2_output_0 : Long(requires_grad=0, device=cpu) = onnx::Constant[value={-1}, onnx_name=\"/Constant_2\"](), scope: __main__.Model:: # /tmp/ipykernel_393401/573755285.py:6:0\n",
      "  %/Mul_output_0 : Long(4, strides=[1], device=cpu) = onnx::Mul[onnx_name=\"/Mul\"](%/ConstantOfShape_output_0, %/Constant_2_output_0), scope: __main__.Model:: # /tmp/ipykernel_393401/573755285.py:6:0\n",
      "  %/Constant_3_output_0 : Long(4, strides=[1], requires_grad=0, device=cpu) = onnx::Constant[value= 1  3 -1  7 [ CPULongType{4} ], onnx_name=\"/Constant_3\"](), scope: __main__.Model:: # /tmp/ipykernel_393401/573755285.py:6:0\n",
      "  %/Equal_output_0 : Bool(4, strides=[1], device=cpu) = onnx::Equal[onnx_name=\"/Equal\"](%/Constant_3_output_0, %/Mul_output_0), scope: __main__.Model:: # /tmp/ipykernel_393401/573755285.py:6:0\n",
      "  %/Where_output_0 : Long(4, strides=[1], device=cpu) = onnx::Where[onnx_name=\"/Where\"](%/Equal_output_0, %/ConstantOfShape_output_0, %/Constant_output_0), scope: __main__.Model:: # /tmp/ipykernel_393401/573755285.py:6:0\n",
      "  %output : Float(1, 3, *, 7, strides=[0, 0, 7, 1], requires_grad=0, device=cpu) = onnx::Expand[onnx_name=\"/Expand\"](%data, %/Where_output_0), scope: __main__.Model:: # /tmp/ipykernel_393401/573755285.py:6:0\n",
      "  return (%output)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "\n",
    "class Model(nn.Module):\n",
    "    def forward(self, x):\n",
    "        return x.expand(1, 3, 4)\n",
    "\n",
    "shape = 1, 1, 4\n",
    "x = torch.rand(*shape)\n",
    "# torch_out = x.expand(1, 3, 4)\n",
    "\n",
    "torch_model = Model()\n",
    "# 导出模型\n",
    "output_name = \"expand\"\n",
    "torch.onnx.export(\n",
    "    torch_model,               # torch 模型\n",
    "    x,                         # 模型输入或者对于多个输入，使用元组\n",
    "    f\"{temp_dir}/{output_name}.onnx\",               # 模型保存的位置（可以是文件或类似文件的对象）\n",
    "    export_params=True,        # 将训练后的参数权重存储在模型文件内\n",
    "    opset_version=17,          # 导出模型的 ONNX 版本\n",
    "    verbose=True,\n",
    "    do_constant_folding=True,  # 是否执行常量折叠以进行优化\n",
    "    input_names = ['data'],    # 模型的输入名称\n",
    "    output_names = ['output'], # 模型的输出名称\n",
    "    dynamic_axes={'data' : {0 : 'batch_size'},    # 可变长度的轴\n",
    "                  'output' : {0 : 'batch_size'}}\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[0.3844, 0.2416],\n",
      "         [0.1347, 0.9315],\n",
      "         [0.4349, 0.9622],\n",
      "         [0.6509, 0.4246]]])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([[[[0.3844, 0.2416],\n",
       "          [0.1347, 0.9315],\n",
       "          [0.4349, 0.9622],\n",
       "          [0.6509, 0.4246]]]])"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = torch.rand(1, 4, 2)\n",
    "print(x)\n",
    "x.expand((1, -1, 4, 2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![模型结构](images/expand.jpg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# onnx_program = torch.onnx.dynamo_export(torch_model, x)\n",
    "# onnx_program.save(\"test_dynamo.onnx\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import onnx\n",
    "import tvm\n",
    "from tvm import relay\n",
    "onnx_model = onnx.load(f\"{temp_dir}/{output_name}.onnx\")\n",
    "mod, params = relay.frontend.from_onnx(onnx_model, {\"data\": shape}, freeze_params=True)\n",
    "# with tvm.transform.PassContext(opt_level=3):\n",
    "#     mod = relay.quantize.prerequisite_optimize(mod, params)\n",
    "mod.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with tvm.transform.PassContext(opt_level=3):\n",
    "    with relay.quantize.qconfig(\n",
    "        skip_conv_layers=[],\n",
    "        # calibrate_mode=\"kl_divergence\", \n",
    "        weight_scale=\"max\",\n",
    "        # round_for_shift=True,\n",
    "        # rounding=\"TONEAREST\", # \"UPWARD\" or \"TONEAREST\"\n",
    "        # calibrate_skip_layers=[],\n",
    "        skip_dense_layer=False,\n",
    "    ):\n",
    "        qmod = relay.quantize.quantize(mod, params)\n",
    "qmod.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "path = \"/media/pc/data/board/arria10/lxw/tasks/tools/npu_user_demos/models/telecom/vehile_det_traffic_yolov8n_c3/yolov8n-c3_384_640_.onnx\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import onnx\n",
    "import tvm\n",
    "from tvm import relay\n",
    "onnx_model = onnx.load(path)\n",
    "mod, params = relay.frontend.from_onnx(onnx_model, {\"images\": (1, 3, 384, 640)}, freeze_params=True)\n",
    "# with tvm.transform.PassContext(opt_level=3):\n",
    "#     mod = relay.quantize.prerequisite_optimize(mod, params)\n",
    "mod.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "xin",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
